//
//  TFCollectionViewColumnLayout.swift
//
//  Created by Andney Jack on 2021/6/21.
//

import UIKit

public class TFCollectionViewColumnLayout: UICollectionViewLayout {
    public var columnCount: Int = 2 {
        willSet {
            if columnCount != newValue {
                invalidateLayout()
            }
        }
    }

    public var columnSpacing: CGFloat = 10.0 {
        willSet {
            if columnSpacing != newValue {
                invalidateLayout()
            }
        }
    }

    public var interitemSpacing: CGFloat = 10.0 {
        willSet {
            if interitemSpacing != newValue {
                invalidateLayout()
            }
        }
    }

    public var headerHeight: CGFloat = 0.0 {
        willSet {
            if headerHeight != newValue {
                invalidateLayout()
            }
        }
    }

    public var footerHeight: CGFloat = 0.0 {
        willSet {
            if footerHeight != newValue {
                invalidateLayout()
            }
        }
    }

    public var headerInset: UIEdgeInsets = .zero {
        willSet {
            if headerInset != newValue {
                invalidateLayout()
            }
        }
    }

    public var footerInset: UIEdgeInsets = .zero {
        willSet {
            if footerInset != newValue {
                invalidateLayout()
            }
        }
    }

    public var sectionInset: UIEdgeInsets = .zero {
        willSet {
            if sectionInset != newValue {
                invalidateLayout()
            }
        }
    }

    public var itemRenderDirection: TFCollectionViewItemRenderDirection = .shortestFirst {
        willSet {
            if itemRenderDirection != newValue {
                invalidateLayout()
            }
        }
    }

    public var allowUseSafeAreaInsets: Bool = true {
        willSet {
            if allowUseSafeAreaInsets != newValue {
                invalidateLayout()
            }
        }
    }

    // TODO: - Pinned Header

    public var sectionHeadersPinToVisibleBounds: Bool = false {
        willSet {
            if sectionHeadersPinToVisibleBounds != newValue {
                invalidateLayout()
            }
        }
    }

    public var sectionFootersPinToVisibleBounds: Bool = false {
        willSet {
            if sectionFootersPinToVisibleBounds != newValue {
                invalidateLayout()
            }
        }
    }

    private var haveFloatingHeaderOrFooter: Bool = false

    // The minimum height of the collection view's content.
    // This could be used to allow hidden headers with no content.
    public var minimumContentHeight: CGFloat = 0.0

    private var delegate: TFCollectionViewDelegateColumnLayout? {
        return collectionView?.delegate as? TFCollectionViewDelegateColumnLayout
    }

    private var sectionInfos = [TFCollectionViewSectionInfo]()

    // MARK: - Public

    public func columnCountForSection(_ section: Int) -> Int {
        return m_columnCountForSection(section)
    }

    public func contentWidthForSection(_ section: Int) -> CGFloat {
        guard let collectionView = collectionView else { return 0.0 }

        let sectionInset = m_insetForSection(at: section)
        var contentWidth = collectionView.bounds.width - sectionInset.left - sectionInset.right

        if allowUseSafeAreaInsets {
            if #available(iOS 11.0, *) {
                let safeAreaInsets = collectionView.safeAreaInsets
                contentWidth -= safeAreaInsets.left + safeAreaInsets.right
            }
        }

        return contentWidth
    }

    public func itemWidthForSection(_ section: Int) -> CGFloat {
        let contentWidth = contentWidthForSection(section)
        let columnCount = m_columnCountForSection(section)
        let columnSpacing = m_columnSpacingForSection(at: section)

        return type(of: self).floorCGFloat((contentWidth - CGFloat(columnCount - 1) * columnSpacing) / CGFloat(columnCount))
    }

    // MARK: - Private Accessors

    private func m_columnCountForSection(_ section: Int) -> Int {
        guard let collectionView = collectionView else { return 0 }

        if delegate?.responds(to: #selector(TFCollectionViewDelegateColumnLayout.collectionView(_:layout:columnCountForSection:))) ?? false {
            return delegate?.collectionView?(collectionView, layout: self, columnCountForSection: section) ?? 0
        } else {
            return columnCount
        }
    }

    private func m_heightForHeaderInSection(_ section: Int) -> CGFloat {
        guard let collectionView = collectionView else { return 0.0 }

        if delegate?.responds(to: #selector(TFCollectionViewDelegateColumnLayout.collectionView(_:layout:heightForHeaderInSection:))) ?? false {
            return delegate?.collectionView?(collectionView, layout: self, heightForHeaderInSection: section) ?? 0.0
        } else {
            return headerHeight
        }
    }

    private func m_itemHeightForSection(_ section: Int, item: Int, itemWidth: CGFloat) -> CGFloat {
        guard let collectionView = collectionView else { return 0.0 }

        if let itemSize = delegate?.collectionView(collectionView, layout: self, sizeForItemAt: IndexPath(item: item, section: section), columnWidth: itemWidth) {
            if itemSize.width > 0 && itemSize.height > 0 {
                return type(of: self).floorCGFloat(itemSize.height * itemWidth / itemSize.width)
            }
        }

        return 0.0
    }

    private func m_heightForFooterInSection(_ section: Int) -> CGFloat {
        guard let collectionView = collectionView else { return 0.0 }

        if delegate?.responds(to: #selector(TFCollectionViewDelegateColumnLayout.collectionView(_:layout:heightForFooterInSection:))) ?? false {
            return delegate?.collectionView?(collectionView, layout: self, heightForFooterInSection: section) ?? 0.0
        } else {
            return footerHeight
        }
    }

    private func m_insetForSection(at section: Int) -> UIEdgeInsets {
        guard let collectionView = collectionView else { return UIEdgeInsets.zero }

        if delegate?.responds(to: #selector(TFCollectionViewDelegateColumnLayout.collectionView(_:layout:insetForSectionAt:))) ?? false {
            return delegate?.collectionView?(collectionView, layout: self, insetForSectionAt: section) ?? UIEdgeInsets.zero
        } else {
            return sectionInset
        }
    }

    private func m_insetForHeaderInSection(_ section: Int) -> UIEdgeInsets {
        guard let collectionView = collectionView else { return UIEdgeInsets.zero }

        if delegate?.responds(to: #selector(TFCollectionViewDelegateColumnLayout.collectionView(_:layout:insetForHeaderInSection:))) ?? false {
            return delegate?.collectionView?(collectionView, layout: self, insetForHeaderInSection: section) ?? UIEdgeInsets.zero
        } else {
            return headerInset
        }
    }

    private func m_insetForFooterInSection(_ section: Int) -> UIEdgeInsets {
        guard let collectionView = collectionView else { return UIEdgeInsets.zero }

        if delegate?.responds(to: #selector(TFCollectionViewDelegateColumnLayout.collectionView(_:layout:insetForFooterInSection:))) ?? false {
            return delegate?.collectionView?(collectionView, layout: self, insetForFooterInSection: section) ?? UIEdgeInsets.zero
        } else {
            return footerInset
        }
    }

    private func m_interitemSpacingForSection(at section: Int) -> CGFloat {
        guard let collectionView = collectionView else { return 0.0 }

        if delegate?.responds(to: #selector(TFCollectionViewDelegateColumnLayout.collectionView(_:layout:interitemSpacingForSectionAt:))) ?? false {
            return delegate?.collectionView?(collectionView, layout: self, interitemSpacingForSectionAt: section) ?? 0.0
        } else {
            return interitemSpacing
        }
    }

    private func m_columnSpacingForSection(at section: Int) -> CGFloat {
        guard let collectionView = collectionView else { return 0.0 }

        if delegate?.responds(to: #selector(TFCollectionViewDelegateColumnLayout.collectionView(_:layout:columnSpacingForSectionAt:))) ?? false {
            return delegate?.collectionView?(collectionView, layout: self, columnSpacingForSectionAt: section) ?? 0.0
        } else {
            return columnSpacing
        }
    }

    private func m_itemRenderDirectionForSection(at section: Int) -> TFCollectionViewItemRenderDirection {
        guard let collectionView = collectionView else {
            return .shortestFirst
        }

        if delegate?.responds(to: #selector(TFCollectionViewDelegateColumnLayout.collectionView(_:layout:itemRenderDirectionForSectionAt:))) ?? false {
            return delegate?.collectionView?(collectionView, layout: self, itemRenderDirectionForSectionAt: section) ?? .shortestFirst
        } else {
            return itemRenderDirection
        }
    }

    private func m_shouldPinToVisibleBounds(ofKind kind: String, section: Int) -> Bool {
        guard let collectionView = collectionView else { return false }

        if kind == TFCollectionElementKindSectionHeader {
            if let ret = delegate?.collectionView?(collectionView, layout: self, shouldPinToVisibleBoundsForSupplementaryViewOfKind: kind, section: section) {
                return ret
            }
            return sectionHeadersPinToVisibleBounds
        } else if kind == TFCollectionElementKindSectionFooter {
            if let ret = delegate?.collectionView?(collectionView, layout: self, shouldPinToVisibleBoundsForSupplementaryViewOfKind: kind, section: section) {
                return ret
            }
            return sectionFootersPinToVisibleBounds
        }
        return false
    }

    // handle floating header or footer
    private func m_layoutAttributesForSupplementaryView(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        guard let collectionView = collectionView, haveFloatingHeaderOrFooter else { return nil }

        var attributeList = [UICollectionViewLayoutAttributes]()

        var allFloatingSectionIndexs = [Int]()
        for section in 0 ..< sectionInfos.count {
            let sectionInfo = sectionInfos[section]

            if sectionInfo.hasFloatingHeaderOrFooter {
                allFloatingSectionIndexs.append(section)
            } else {
                // normal header footer
                if let attributes = sectionInfo.headerAttribute, sectionInfo.headerHeight > 0 {
                    var headerFrame = attributes.frame
                    headerFrame.origin.y = sectionInfo.headerBounds.top

                    if headerFrame.intersects(rect) {
                        attributeList.append(attributes)
                    }
                }

                if let attributes = sectionInfo.footerAttribute, sectionInfo.footerHeight > 0 {
                    var footerFrame = attributes.frame
                    footerFrame.origin.y = sectionInfo.footerBounds.top

                    if footerFrame.intersects(rect) {
                        attributeList.append(attributes)
                    }
                }
            }
        }

        guard allFloatingSectionIndexs.count > 0 else {
            return attributeList
        }

        let offsetY = collectionView.contentOffset.y
        let top = offsetY

        var lastFloatingItemTop: CGFloat?
        var flag: Bool = false
        var skip: Bool = false
        for idx in (0 ..< allFloatingSectionIndexs.count).reversed() {
            let section = allFloatingSectionIndexs[idx]
            let sectionInfo = sectionInfos[section]

            if sectionInfo.headerBounds.top >= top || skip {
                if let attributes = sectionInfo.headerAttribute {
                    type(of: self).updateLayoutAttributesTop(sectionInfo.headerBounds.top, for: attributes)
                    attributes.zIndex = 0
                    attributeList.append(attributes)

                    lastFloatingItemTop = attributes.frame.minY
                }
            } else {
                if let attributes = sectionInfo.headerAttribute {
                    var top: CGFloat = max(0, offsetY)
                    if let lastTop = lastFloatingItemTop, top + sectionInfo.headerHeight > lastTop {
                        top = lastTop - sectionInfo.headerHeight
                        flag = true

                        type(of: self).updateLayoutAttributesTop(top, for: attributes)
                        attributes.zIndex = 1024 + section
                        attributeList.append(attributes)

                    } else {
                        if flag {
                            type(of: self).updateLayoutAttributesTop(sectionInfo.headerBounds.top, for: attributes)
                            attributes.zIndex = 0
                            attributeList.append(attributes)

                        } else {
                            type(of: self).updateLayoutAttributesTop(top, for: attributes)
                            attributes.zIndex = 1024 + section
                            attributeList.append(attributes)
                        }

                        skip = true
                    }
                }
            }
        }
        return attributeList
    }

    // MARK: - Override Methods

    override public func prepare() {
        super.prepare()

        sectionInfos.removeAll()
        haveFloatingHeaderOrFooter = false

        guard let collectionView = collectionView, collectionView.numberOfSections > 0 else {
            return
        }

        assert(delegate != nil, "UICollectionView's delegate should conform to CTCollectionViewDelegateColumnLayout protocol")
        assert(columnCount > 0 || (delegate?.responds(to: #selector(TFCollectionViewDelegateColumnLayout.collectionView(_:layout:columnCountForSection:))) ?? false),
               "CTCollectionViewColumnLayout's columnCount should be greater than 0, or delegate must implement -columnCountForSection:")

        var top: CGFloat = 0.0

        var safeAreaInsets = UIEdgeInsets.zero
        if allowUseSafeAreaInsets {
            if #available(iOS 11.0, *) {
                safeAreaInsets = collectionView.safeAreaInsets
            }
        }

        for section in 0 ..< collectionView.numberOfSections {
            let sectionInfo = TFCollectionViewSectionInfo()
            sectionInfo.section = section
            sectionInfo.columnCount = m_columnCountForSection(section)

            sectionInfo.columnSpacing = m_columnSpacingForSection(at: section)
            sectionInfo.interitemSpacing = m_interitemSpacingForSection(at: section)
            sectionInfo.collectionViewWidth = collectionView.bounds.width
            sectionInfo.columnWidth = itemWidthForSection(section)

            sectionInfo.top = top
            sectionInfo.headerHeight = m_heightForHeaderInSection(section)
            sectionInfo.footerHeight = m_heightForFooterInSection(section)
            sectionInfo.headerInsets = m_insetForHeaderInSection(section)
            sectionInfo.footerInsets = m_insetForFooterInSection(section)
            sectionInfo.sectionInsets = m_insetForSection(at: section)
            sectionInfo.safeAreaInsets = safeAreaInsets
            sectionInfo.itemRenderDirection = m_itemRenderDirectionForSection(at: section)

            if sectionInfo.headerHeight > 0 {
                sectionInfo.headerPinToVisible = m_shouldPinToVisibleBounds(ofKind: TFCollectionElementKindSectionHeader, section: section)
            } else {
                sectionInfo.headerPinToVisible = false
            }
            if sectionInfo.footerHeight > 0 {
                sectionInfo.footerPinToVisible = m_shouldPinToVisibleBounds(ofKind: TFCollectionElementKindSectionFooter, section: section)
            } else {
                sectionInfo.footerPinToVisible = false
            }

            if !haveFloatingHeaderOrFooter && sectionInfo.hasFloatingHeaderOrFooter {
                haveFloatingHeaderOrFooter = true
            }

            let itemCount = collectionView.numberOfItems(inSection: section)
            for idx in 0 ..< itemCount {
                let itemInfo = TFCollectionViewItemInfo()
                itemInfo.section = section
                itemInfo.item = idx
                itemInfo.itemHeight = m_itemHeightForSection(section, item: idx, itemWidth: sectionInfo.columnWidth)

                sectionInfo.allItemInfos.append(itemInfo)
            }

            // invoke this method will calculate all layout infos.
            top += sectionInfo.totalSectionHeight
            sectionInfos.append(sectionInfo)
        }
    }

    override public var collectionViewContentSize: CGSize {
        guard let collectionView = self.collectionView, collectionView.numberOfSections > 0 else {
            return CGSize.zero
        }

        var contentSize = CGSize(width: collectionView.bounds.width, height: 0.0)
        if let lastSectionInfo = sectionInfos.last {
            contentSize.height = lastSectionInfo.sectionBottom
        }

        if contentSize.height < minimumContentHeight {
            contentSize.height = minimumContentHeight
        }
        return contentSize
    }

    override public func layoutAttributesForItem(at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        if indexPath.section >= 0 && indexPath.section < sectionInfos.count {
            let sectionInfo = sectionInfos[indexPath.section]
            if indexPath.item >= 0 && indexPath.item < sectionInfo.allItemInfos.count {
                return sectionInfo.allItemInfos[indexPath.item].itemAttributes
            }
        }
        return nil
    }

    override public func layoutAttributesForSupplementaryView(ofKind elementKind: String, at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        if indexPath.section >= 0 && indexPath.section < sectionInfos.count {
            let sectionInfo = sectionInfos[indexPath.section]

            if elementKind == TFCollectionElementKindSectionHeader {
                type(of: self).updateLayoutAttributesTop(sectionInfo.headerBounds.top, for: sectionInfo.headerAttribute)
                return sectionInfo.headerAttribute
            } else if elementKind == TFCollectionElementKindSectionFooter {
                type(of: self).updateLayoutAttributesTop(sectionInfo.footerBounds.top, for: sectionInfo.footerAttribute)
                return sectionInfo.footerAttribute
            }
        }

        return nil
    }

    override public func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        var attributesList = [UICollectionViewLayoutAttributes]()

        var stop = false
        var exists = false // 是否已发现了交叉区域
        for sectionInfo in sectionInfos {
            if stop {
                break
            }

            let hasHeader = sectionInfo.headerHeight > 0.0 ? true : false
            let hasFooter = sectionInfo.footerHeight > 0.0 ? true : false

            if !exists {
                if let footerAttribute = sectionInfo.footerAttribute {
                    if hasFooter && rect.minY > footerAttribute.frame.maxY {
                        continue
                    }
                }
            } else {
                if let headerAttribute = sectionInfo.headerAttribute {
                    if hasHeader && rect.maxY < headerAttribute.frame.minY {
                        stop = true
                        break
                    }
                }
            }

            if let headerAttribute = sectionInfo.headerAttribute {
                var headerFrame = headerAttribute.frame
                headerFrame.origin.y = sectionInfo.headerBounds.top

                if hasHeader && headerFrame.intersects(rect) {
                    exists = true
                    if !haveFloatingHeaderOrFooter {
                        headerAttribute.frame = headerFrame
                        attributesList.append(headerAttribute)
                    }
                }
            }

            if let footerAttribute = sectionInfo.footerAttribute {
                var footerFrame = footerAttribute.frame
                footerFrame.origin.y = sectionInfo.footerBounds.top

                if hasFooter && footerFrame.intersects(rect) {
                    exists = true
                    if !haveFloatingHeaderOrFooter {
                        footerAttribute.frame = footerFrame
                        attributesList.append(footerAttribute)
                    }
                }
            }

            for itemInfo in sectionInfo.allItemInfos {
                if let itemAttributes = itemInfo.itemAttributes {
                    if itemAttributes.frame.intersects(rect) {
                        exists = true
                        attributesList.append(itemAttributes)
                    }
                }
            }
        }

        if haveFloatingHeaderOrFooter {
            if let headerFooterAttributesList = m_layoutAttributesForSupplementaryView(in: rect),
               headerFooterAttributesList.count > 0 {
                attributesList.append(contentsOf: headerFooterAttributesList)
            }
        }

        return attributesList
    }

    override public func shouldInvalidateLayout(forBoundsChange newBounds: CGRect) -> Bool {
        guard let collectionView = collectionView else {
            return super.shouldInvalidateLayout(forBoundsChange: newBounds)
        }

        if haveFloatingHeaderOrFooter {
            return true
        }

        if collectionView.bounds.width != newBounds.width {
            return true
        }
        return false
    }

    override public class var layoutAttributesClass: AnyClass {
        return TFCollectionViewColumnLayout.self
    }
}

// MARK: - Util

extension TFCollectionViewColumnLayout {
    fileprivate static func floorCGFloat(_ value: CGFloat) -> CGFloat {
        let scale = UIScreen.main.scale
        return floor(value * scale) / scale
    }

    // update top of the frame for layout attributes
    fileprivate static func updateLayoutAttributesTop(_ top: CGFloat, for layoutAttributes: UICollectionViewLayoutAttributes?) {
        guard let attributes = layoutAttributes else { return }

        if attributes.frame.minY != top {
            var frame = attributes.frame
            frame.origin.y = top
            attributes.frame = frame
        }
    }
}

fileprivate struct TFSectionInfo {
    var section: Int
    var haveHeader: Bool
    var haveFooter: Bool

    init(section: Int, haveHeader: Bool, haveFooter: Bool) {
        self.section = section
        self.haveHeader = haveHeader
        self.haveFooter = haveFooter
    }
}
